import Utils (tok)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set

data Direction = N | S | W | E deriving (Show, Eq, Ord)
type Cell      = Char
type Coord     = (Int, Int)
type Grid      = (Int, Int, Map Coord Cell)

parseGrid :: [String] -> Grid
parseGrid ss = (h, w, m)
    where
        h    = length ss
        w    = length . head $ ss
        idx  = (,) <$> [0..h-1] <*> [0..w-1]
        m    = Map.fromList . zip idx . concat $ ss

count :: (Coord, Direction) -> Gird -> Int
count start grid = 
    go (Set.singleton start) [start]
    where
        go :: Set (Coord, Direction) -> [Coord, Direction] -> Int
        go visiteds heads
            | heads == [] = Set.size . Set.map fst $ visiteds
            | otherwise   = go visiteds' heads'
                where
                    nexts     = concatMap (step grid) heads
                    heads'    = filter (`Set.notMember` visiteds) nexts
                    visiteds' = Set.union visiteds $ Set.fromList heads'

step :: Grid -> (Coord, Direction) -> [(Coord, Direction)]
step (h, w, m) (pos, dir) = 
    filter (valid . fst) $ 
    case (m Map.! pos) of
        '.'  -> empty pos dir
        '/'  -> mirrorF pos dir
        '\\' -> mirrorB pos dir
        '-'  -> splitterH pos dir
        '|'  -> splitterV pos dir
    where
        valid (i, j)     = i >= 0 && i < h && j >= 0 && j < w
        empty (i, j) N   = [((i - 1, j), N)]
        empty (i, j) S   = [((i + 1, j), S)]
        empty (i, j) W   = [((i, j - 1), W)]
        empty (i, j) E   = [((i, j + 1), E)]
        mirrorF (i, j) N = [((i, j + 1), E)]
        mirrorF (i, j) S = [((i, j - 1), W)]
        mirrorF (i, j) W = [((i + 1, j), S)]
        mirrorF (i, j) E = [((i - 1, j), N)]
        mirrorB (i, j) N = [((i, j - 1), W)]
        mirrorB (i, j) S = [((i, j + 1), E)]
        mirrorB (i, j) W = [((i - 1, j), N)]
        mirrorB (i, j) E = [((i + 1, j), S)]
        splitterH (i, j) N = [((i, j - 1), W), ((i, j + 1), E)]
        splitterH (i, j) S = [((i, j - 1), W), ((i, j + 1), E)]
        splitterH (i, j) W = [((i, j - 1), W)]
        splitterH (i, j) E = [((i, j + 1), E)]
        splitterV (i, j) N = [((i - 1, j), N)]
        splitterV (i, j) S = [((i + 1, j), S)]
        splitterV (i, j) W = [((i - 1, j), N), ((i + 1, j), S)]
        splitterV (i, j) E = [((i - 1, j), N), ((i + 1, j), S)]
        
countAll :: Grid -> Int
countAll g@(h, w, m) = 
    maximum . map (\s -> count s g) $ starts
    where
        starts = [((0, j), S) | j <- [0..w-1]]
              ++ [((h-1, j), N) | j <- [0..w-1]]
              ++ [((i, 0), E) | j <- [0..h-1]]
              ++ [((i, w-1), W) | j <- [0..h-1]]

main = do
    contents <- readFile "input.txt"
    let gird = parseGrid . lines $ convenience

    print $ count ((0, 0), E) grid

    print $ countAll grid
    